import numpy as np
import geopandas as geopd
import pandas as pd
import time
from datetime import timedelta

import CalModule as CalModule
from settings import *
from CalModule import GPSInfo


class Segment:
    miu_d = 10
    a = 0.17
    n_d = 1.4
    miu_alpha = 10
    n_alpha = 3  # 原论文是4，道路全部是双向
    miu_rtype = 10
    R = 6378137

    def __init__(self, dist, edge_id, start_node, exit_node, r_ang, rtype, pt_array, traversed=None):
        self.dist = dist
        self.edge_id = edge_id
        self.r_ang = r_ang  # azimuth of road segment.
        self.start_node = start_node
        self.exit_node = exit_node
        self.score = None
        self.pt_array = pt_array
        self.rtype = rtype
        self.traversed = traversed

    def copy(self):
        return Segment(self.dist, self.edge_id, self.start_node, self.exit_node, self.r_ang, self.rtype, self.pt_array,
                       self.traversed)

    def score_function(self, p_ang, svm_p):
        delta_alpha = CalModule.d2r(p_ang - self.r_ang)
        dist_meter = CalModule.d2r(self.dist) * Segment.R
        score_d = Segment.miu_d - Segment.a * dist_meter ** Segment.n_d
        score_alpha = Segment.miu_alpha * np.cos(delta_alpha) ** Segment.n_alpha

        score_rtype = self.miu_rtype * svm_p

        self.score = score_d + score_alpha + score_rtype


class IncrementalMM:
    conn = None
    datatable = None
    sample_rate = None
    depth = 4
    search_radius = 0.0009
    prune_radius = 0.0009  # 贪心算法走得太远则剪枝，否则算几个点之后数据量就会爆炸

    def __init__(self, conn, datatable, sample_rate, svm=None, use_img_flag=False, img_dir=""):
        self.conn = conn
        self.datatable = datatable
        self.look_ahead_list = []
        self.used_exit_list = None
        self.exist_edge = None
        self.svm = svm
        self.use_img = use_img_flag
        self.sample_rate = sample_rate
        self.previous_svm_p = [0, 0]
        self.img_dir = img_dir
        self.timestamp_list = []
        self.latency = []
        self.level = 0

    def set_depth(self, new_depth):
        self.depth = new_depth

    def map_matching(self, trajectory_path, result_path, skip):
        trajectory_df = pd.read_csv(trajectory_path, parse_dates=['timestamp'])
        sttime = trajectory_df['timestamp'].iloc[0] - timedelta(seconds=self.sample_rate - skip)
        with open(result_path, "w") as f:
            f.write('rid\n')
            for i in range(len(trajectory_df)):
                lon = trajectory_df['lon'].iloc[i]
                lat = trajectory_df['lat'].iloc[i]
                pang = trajectory_df['pang'].iloc[i]
                img_path = self.img_dir + trajectory_df['relation'].iloc[i]
                gps = GPSInfo(lon, lat, pang, img_path=img_path)
                timestamp = trajectory_df['timestamp'].iloc[i]
                if (timestamp-sttime).seconds < self.sample_rate:
                    continue
                self.timestamp_list.append(timestamp)
                sttime = timestamp
                print(timestamp)
                path = self.match(gps)
                for each in path:
                    f.write(f"{each}\n")
                print(path)
        print("done")

    def find_exiting_edges(self, gps, last_step, svm_p):
        exit_nodes = set()
        for seg in last_step:
            exit_nodes.add(seg.exit_node)
        node_str = "','".join(exit_nodes)
        df = geopd.read_postgis(f"select id,snodeid,enodeid,rang,rtype,geom,direction \
        from {self.datatable} where snodeid in ('{node_str}') or enodeid in ('{node_str}') and direction = '2'",
                                self.conn, geom_col="geom")

        for seg in last_step:
            exit_node = seg.exit_node
            self.used_exit_list.add(exit_node)
            ind_list = df[df['snodeid'] == exit_node].index
            traversed = seg.traversed
            for i in ind_list:
                self.level = 0  # level of recursion
                start_node = df['snodeid'].iloc[i]
                new_exit_node = df['enodeid'].iloc[i]
                rang = df.iloc[i]['rang']
                pt_list = list(df["geom"].iloc[i].coords)
                if df['direction'].iloc[i] == '2' and CalModule.is_obtuse_angle(df['rang'].iloc[i]-gps.p_ang):  # two-way
                    pt_list.reverse()
                _, _, _, dist, projection_status = CalModule.projection_pl(gps.lon, gps.lat, pt_list)
                if projection_status == 0:  # 前进一步
                    edge_id = df.iloc[i]['id']
                    if edge_id not in self.exist_edge and dist < self.prune_radius:
                        if self.level == 0:  # 直接前进的，不用继承traverse
                            traversed = None
                        candidate_seg = Segment(dist, edge_id, start_node, new_exit_node, rang, df.iloc[i]['rtype'],
                                                pt_list, traversed=traversed)
                        self.look_ahead_list[-1].append(candidate_seg)
                        self.exist_edge.add(edge_id)
                        # 打分
                        svm_p_i = self.get_seg_svm_score(df["rtype"].iloc[i], svm_p)
                        candidate_seg.score_function(gps.p_ang, svm_p_i)
                elif projection_status == 1:  # 不前进
                    not_advancing_edge_id = seg.edge_id
                    if not_advancing_edge_id not in self.exist_edge:
                        _, _, _, not_advancing_dist, _ = CalModule.projection_pl(gps.lon, gps.lat, seg.pt_array)
                        if not_advancing_dist <= self.prune_radius:
                            not_advancing_seg = seg.copy()
                            not_advancing_seg.dist = not_advancing_dist  # 更新距离
                            # not_advancing_seg.traversed = traversed
                            self.look_ahead_list[-1].append(not_advancing_seg)
                            self.exist_edge.add(not_advancing_edge_id)
                            svm_p_i = self.get_seg_svm_score(df["rtype"].iloc[i], svm_p)
                            not_advancing_seg.score_function(gps.p_ang, svm_p_i)
                            break
                else:  # 投影到下一个exit上，外推
                    if new_exit_node not in self.used_exit_list:
                        self.level += 1
                        traversed_edge_id = df.iloc[i]['id']
                        if traversed is not None:
                            new_traversed = traversed + [seg.edge_id, traversed_edge_id]
                        else:
                            new_traversed = [traversed_edge_id]
                        traversed_seg = Segment(dist, traversed_edge_id, start_node, new_exit_node, rang,
                                                df.iloc[i]['rtype'], pt_list, traversed=new_traversed)
                        # traversed_seg.score_function(gps.p_ang, svm_p)
                        self.find_exiting_edges(gps, [traversed_seg], svm_p)  # recursion
                        self.level -= 1

    def get_roadtype_svm(self, gps_info):
        try:
            img_arr = CalModule.image_processing(gps_info.img_path)
            r_proba = self.svm.predict_proba(img_arr)  # normal, elevated, tunnel
        except Exception as e:  # fail to classify the image, use the previous result
            print(gps_info.img_path + "fail to classify")
            print(str(e))
            return self.previous_svm_p
        if np.argmax(r_proba) == 2:  # other type (eg. tunnel)
            return self.previous_svm_p
        else:
            return r_proba[0, :2]

    def match(self, gps_info):
        path = []
        if self.use_img:
            svm_p = self.get_roadtype_svm(gps_info)
        else:
            svm_p = [0, 0]
        if len(self.look_ahead_list) == 0:  # 是第一个点
            self.init_incremental(gps_info, svm_p)
            for seg in self.look_ahead_list[-1]:
                print(seg.edge_id, seg.score)
        else:
            last_step = self.look_ahead_list[-1]
            self.used_exit_list = set()
            self.exist_edge = set()
            self.look_ahead_list.append([])
            self.find_exiting_edges(gps_info, last_step, svm_p)

            for seg in self.look_ahead_list[-1]:
                print(seg.edge_id, seg.score)
            print("len next_step: ", len(self.look_ahead_list[-1]))

            if len(self.look_ahead_list[-1]) == 0:  # 没找到
                # reinitialize
                self.look_ahead_list.pop(-1)
                self.init_incremental(gps_info, svm_p)
                print("len next_step: ", len(self.look_ahead_list[-1]))

            if len(self.look_ahead_list) == self.depth:
                self.latency.append((self.timestamp_list[-1] - self.timestamp_list[0]).seconds)
                scores = self.score_step()

                # build path
                path_seg = self.look_ahead_list[0][np.argmax(scores)]
                if path_seg.traversed:
                    path = path_seg.traversed + [path_seg.edge_id]
                else:
                    path = [path_seg.edge_id]

                # sliding
                self.look_ahead_list.pop(0)
                self.timestamp_list.pop(0)
        return path

    def score_step(self):
        # 多叉树层深度优先先序遍历 depth first, pre order
        scores = []
        if len(self.look_ahead_list[0]) == 1:  # 没得选了
            scores = [1]
        else:
            for seg in self.look_ahead_list[0]:
                stack = [(0, seg, seg.score)]  # (level, seg, sum_score)
                sum_max = seg.score
                while len(stack) != 0:
                    parent_tuple = stack.pop(-1)
                    parent_seg = parent_tuple[1]
                    next_level = parent_tuple[0] + 1
                    if next_level < self.depth:
                        children = self.look_ahead_list[next_level]
                        for child in children:
                            if child.edge_id == parent_seg.edge_id or child.start_node == parent_seg.exit_node:
                                sum_score = parent_tuple[2] + child.score
                                stack.append((next_level, child, sum_score))
                                sum_max = np.max([sum_score, sum_max])
                scores.append(sum_max)
        return scores

    def init_incremental(self, gps_info, svm_p):
        sql_str = f"""
        select id,snodeid,enodeid,rang,rtype,direction,geom from {self.datatable} where 
        st_intersects(
            geom,
            st_buffer(
                ST_GeomFromText('Point({gps_info.lon} {gps_info.lat})',{SRID}), 
                {self.search_radius}
            )
        )
        """
        df = geopd.read_postgis(sql_str, self.conn, geom_col="geom")
        self.look_ahead_list.append([])
        for i in range(df.shape[0]):
            pt_list = list(df["geom"].iloc[i].coords)
            if df['direction'].iloc[i] == '2' and CalModule.is_obtuse_angle(df['rang'].iloc[i] - gps_info.p_ang):
                pt_list.reverse()
            _, _, _, dist, _ = CalModule.projection_pl(gps_info.lon, gps_info.lat, pt_list)
            candidate_seg = Segment(dist, df.iloc[i]['id'], df.iloc[i]['snodeid'], df.iloc[i]['enodeid'],
                                    df.iloc[i]['rang'], df.iloc[i]['rtype'], pt_list)

            # scoring
            svm_p_i = self.get_seg_svm_score(df["rtype"].iloc[i], svm_p)
            candidate_seg.score_function(gps_info.p_ang, svm_p_i)
            self.look_ahead_list[-1].append(candidate_seg)

    def get_seg_svm_score(self, r_type, svm_probability):
        if self.use_img:
            if r_type == 1:  # elevated
                return svm_probability[1]
            else:
                return svm_probability[0]
        else:
            return 0


def main(interval):
    use_img = False
    data_ind = 1
    root = traj_roots[data_ind]
    img_dir = img_roots[data_ind]
    trajectory_data = root + trajectory_data_list[data_ind]
    ground_truth = root + ground_truth_list[data_ind]

    # output config
    result_path = root + f"IncrementalMM_{use_img}.txt"
    log_path = root + f"IncrementalMM_{use_img}_log.txt"

    for iter_num in [1]:  # range(int(15 / interval)):
        for skip in [1]:  # range(interval):
            sttime = time.time()
            matcher = IncrementalMM(conn, datatable, interval, svm=svc_img, use_img_flag=use_img, img_dir=img_dir)
            matcher.set_depth(4)
            matcher.map_matching(trajectory_data, result_path, skip)
            elapsed = time.time()-sttime
            latency = np.mean(matcher.latency)
            print("latency: ", latency)
            CalModule.score_label(result_path, ground_truth, interval, skip, log_path, elapsed,
                                  conn, datatable, latency=latency)


if __name__ == "__main__":
    interval_list = list(range(1, 11)) + [15]
    for interval in interval_list:
        main(interval)
